from debugtools import *
from gputools import *
from pool import *
from baserpc import *
from multiprocessing import Process
from threading import Thread
import time

__all__ = ['DEFAULT_PORT', 'RpcServer', 'RpcClient', 'RpcClients', 
           'open_server', 'close_servers', 'Procedure']

def open_server(port = None):
    RpcServer(port=port).listen()

def close_servers(addresses):
    if type(addresses) is str:
        addresses = [(addresses, DEFAULT_PORT)]
    if type(addresses) is tuple:
        addresses = [addresses]
    RpcClients(addresses).close_server()




class Procedure(object):
    """
    A procedure sent by a client to a server.
    object handler_id, instance of handler_class, call of method_name(*args, **kwds)
    """
    def __init__(self, handler_class, handler_id, name, args, kwds):
        self.handler_class = handler_class
        if handler_class is not None:
            self.handler_class_name = handler_class.__name__
        else:
            # handler_class = None means that the method is a method of the RpcServer
            # not a method of handler_class
            self.handler_class_name = 'native' 
        if args is None:
            args = ()
        if kwds is None:
            kwds = {}
        self.handler_id = handler_id
        self.name = name
        self.args = args
        self.kwds = kwds
    
    def is_native(self):
        return self.handler_class_name == 'native'

    def __repr__(self):
        return "%s.%s" % (self.handler_class_name, self.name)
#        return "<Procedure %s.%s with args %s, kwds %s>" % (self.handler_class_name, self.name, 
#                                                            str(self.args), str(self.kwds))




class RpcServer(BaseRpcServer):
    def initialize(self):
        global MAXGPU
        if CANUSEGPU:
            MAXGPU = initialise_cuda()
        
        self.pools = {}
        self.pools['CPU'] = CustomPool.create_pool(MAXCPU) # when creating a CustomPool, set its 
                                                           # pool property to this if using CPUs
        if CANUSEGPU:
            self.pools['GPU'] = CustomPool.create_pool(MAXGPU)
    
        self.resources = {}
        self.resources['CPU'] = {}
        if CANUSEGPU:
            self.resources['GPU'] = {}
        
        self.handlers = {}
    
    def set_resources(self, client, type, units):
        """
        Allocates units type to client
        type is CPU or GPU
        """
        self.resources[type][client] = units
    
    def get_resources(self, client):
        return self.resources
    
    def add_handler(self, client, handler_class, handler_id):
        log.debug("server: creating new handler '%s'" % handler_id)
        self.handlers[handler_id] = handler_class() # Object initialization
        self.handlers[handler_id].pools = self.pools # the handler objects can access to the global pools
    
    def delete_handler(self, client, handler_id):
        if handler_id is None:
            handler_id = client
        log.debug("server: deleting handler '%s'" % handler_id)
        del self.handlers[handler_id]
    
    def process(self, client, procedure):
        if procedure.is_native():
            # call self.name(*args, **kwds)
            result = getattr(self, procedure.name)(client, *procedure.args, **procedure.kwds)
        else:
            # Default handler id is the client IP
            if procedure.handler_id is None:
                procedure.handler_id = client
            # creates the handler if it doesn't exist
            if procedure.handler_id not in self.handlers.keys():
                self.add_handler(client, procedure.handler_class, procedure.handler_id)
            # call self.handlers[id](*args, **kwds)
            result = getattr(self.handlers[procedure.handler_id], procedure.name)(*procedure.args, **procedure.kwds)
        return result
    
    def shutdown(self):
        for id, handler in self.handlers.iteritems():
            if hasattr(handler, 'close'):
                log.debug("closing handler '%s'" % id)
                handler.close()
        for type, pool in self.pools.iteritems():
            log.debug("closing pool of %s" % type)
            pool.close()




class RpcClient(BaseRpcClient):
    def __init__(self, server, handler_class = None, handler_id = None):
        BaseRpcClient.__init__(self, server)
        self.handler_class = handler_class
        self.handler_id = handler_id

    def execute_method(self, handler_class, handler_id, method, *args, **kwds):
        procedure = Procedure(handler_class, handler_id, method, args, kwds)
        return self.execute(procedure)
    
    def execute_native(self, method, *args, **kwds):
        return self.execute_method(None, None, method, *args, **kwds)

    def set_handler_id(self, handler_id):
        log.debug("client: setting handler id='%s'" % str(handler_id))
        self.handler_id = handler_id

    def add_handler(self, handler_id):
        self.set_handler_id(handler_id)
        log.debug("client: adding handler '%s'" % str(handler_id))
        return self.execute_native('add_handler', self.handler_class, handler_id)

    def delete_handler(self, handler_id = None):
        log.debug("client: deleting handler '%s'" % str(handler_id))
        return self.execute_native('delete_handler', handler_id)

    def set_resources(self, units, type = 'CPU'):
        return self.execute_native('set_resources')
    
    def get_resources(self):
        return self.execute_native('get_resources')

    def __getattr__(self, method):
        log.debug("getting attribute '%s'" % method)
        return lambda *args, **kwds: self.execute_method(self.handler_class, self.handler_id, method, *args, **kwds)




class RpcClients(BaseRpcClients):
    def __init__(self, servers, handler_class = None, handler_id = None):
        self.servers = servers
        self.handler_class = handler_class
        self.handler_id = handler_id
        
        self.clients = [RpcClient(server, handler_class, handler_id) for server in servers]
        self.results = {}
        
    def distribute(self, handler_class, handler_id, name, *argss, **kwdss):
        procedures = []
        k = len(argss) # number of non-named arguments
        keys = kwdss.keys() # keyword arguments
        
        # Duplicates non-list args
        argss = list(argss)
        for l in xrange(k):
            if type(argss[l]) is not list:
                argss[l] = [argss[l]]*len(self.servers)
        
        # Duplicates non-list kwds
        for key in keys:
            if type(kwdss[key]) is not list:
                kwdss[key] = [kwdss[key]]*len(self.servers)
        
        for i in xrange(len(self.servers)):
            args = [argss[l][i] for l in xrange(k)]
            kwds = dict([(key, kwdss[key][i]) for key in keys])
            procedure = Procedure(handler_class, handler_id, name, args, kwds)
            procedures.append(procedure)
        
        return self.execute(procedures)
    
    def set_handler_id(self, handler_id):
        return self.distribute(None, None, 'set_handler_id', [handler_id]*len(self.servers))

    def add_handler(self, handler_id):
        return self.distribute(None, None, 'add_handler', [handler_id]*len(self.servers))

    def delete_handler(self, handler_id = None):
        return self.distribute(None, None, 'delete_handler', [handler_id]*len(self.servers))

    def get_resources(self):
        return self.distribute(None, None, 'get_resources')
    
    def get_resources(self, units, type='CPU'):
        return self.distribute(None, None, 'get_resources', units, type)
    
    def __getattr__(self, name):
        log.debug("getting attribute '%s' on RpcClients" % name)
        return lambda *argss, **kwdss: self.distribute(self.handler_class, self.handler_id, name, *argss, **kwdss)

